
import random

import pymongo
from pymongo import MongoClient
from domain import Proxy
from settings import MONGO_URL
from utils.log import logger


class MongoPool(object):
    def __init__(self):
        """建立数据库连接，获取要操作的集合"""
        self.client = MongoClient(MONGO_URL)
        # 获取要操作的集合
        self.proxies = self.client['proxies_pool']['proxies']

    def __del__(self):
        """关闭数据库"""
        self.client.close()

    def insert_one(self, proxy):
        count = self.proxies.count({'_id': proxy.ip})
        if count == 0:
            dic = proxy.__dict__
            # 使用proxy作为主键
            dic['_id'] = proxy.ip
            self.proxies.insert_one(dic)
            logger.info("插入新的代理:{}".format(proxy))
        else:
            logger.warning("已存在的代理:{}".format(proxy))

    def update_one(self, proxy):
        self.proxies.update_one({'_id': proxy.ip}, {'$set': proxy.__dict__})
        logger.info("已更新的代理:{}".format(proxy))

    def delete_one(self, proxy):
        """根据指定的ip删除代理"""
        self.proxies.delete_one({'_id': proxy.ip})
        logger.info("删除代理ip:{}".format(proxy))

    def find_all(self):
        """查询所有代理ip的功能"""
        cursor = self.proxies.find()

        for item in cursor:
            # 删除 _id 这个key
            item.pop("_id")

            proxy = Proxy(**item)
            yield proxy

    def find(self, condition={}, count=0):
        """
        :description: 实现查询的功能，根据条件进行查询，可以指定查询条件，先分数降序，速度升序排，保证优质的代理IP
        :param condition: 查询条件字典
        :param count: 限制最多取出多少个代理IP
        :return: 返回满足要求的代理IP(Proxy对象)列表
        """
        cursor = self.proxies.find(condition, limit=count).sort([
            ('score', pymongo.DESCENDING), ('speed', pymongo.ASCENDING)
        ])
        # 存储查询处理的IP
        proxy_list = []
        for item in cursor:
            item.pop("_id")
            proxy = Proxy(**item)
            proxy_list.append(proxy)

        return proxy_list

    def get_proxies(self, protocol=None, domain=None, count=0, nick_type=0):
        """
        实现根据协议类型 和要访问网站的域名，获取ip列表
        :param protocol: 协议：http、HTTPS
        :param domain: 域名：baidu.com
        :param count: 用于限制获取多个代理ip，默认是获取所有
        :param nick_type: 匿名类型，默认，获取高匿的代理ip
        :return: 满足要求的代理ip
        """
        # 定义查询条件
        conditions = {'nick_type': nick_type}
        # 根据协议，指定查询条件
        if protocol is None:
            # 如果没有协议, 就获取及支持http 又支持https的协议
            conditions['protocol'] = 2
        elif protocol.lower() == 'http':
            # 如果是HTTP的请求, 使用支持http 和 支持http/https均可以
            conditions['protocol'] = {'$in': [2, 0]}
        elif protocol.lower() == 'https':
            # 如果是HTTP的请求, 使用支持https 和 支持http/https均可以
            conditions["protocol"] = {'$in': [2, 1]}

        if domain:
            conditions['disable_domains'] = {'$nin': [domain]}

        print(conditions)
        return self.find(conditions, count=count)

    def random(self, protocol=None, domain=None, count=0):
        """
        从指定数量代理IP中, 随机获取一个
        :param protocol: 协议: http 或 https
        :param domain: 要访问网站的域名
        :param count: 代理IP的数量
        :return: 一个随机获取代理IP
        """
        proxy_list = self.get_proxies(protocol, domain=domain, count=count)
        return random.choice(proxy_list)

    def random_one(self, count=5):
        """
        从代理池中选择一个响应速度 < 1s的IP
        """
        proxy_list = self.find(condition={'speed': {"$gte": 0, "$lte": 1}}, count=count)
        return random.choice(proxy_list)

    def disable_domain(self, ip, domain):
        # 如果该IP下没有这个域名才更新
        if self.proxies.count_documents({'_id': ip, 'disable_domains': domain}) == 0:
            # 向指定IP的不可用域名中添加域名
            self.proxies.update_one({'_id': ip}, {'$push': {'disable_domains': domain}})


if __name__ == '__main__':
    mongo = MongoPool()

    # 1. insert_one()测试  --> ok
    proxy1 = Proxy('111.111.111.111', port='8008')
    mongo.insert_one(proxy1)

    # 2. update_one() 测试  --> ok
    mongo.update_one(Proxy('111.111.111.111', port='1111', area="贵州省小雅镇"))

    # 3. delete_one() 测试  --> ok
    mongo.delete_one(Proxy('111.111.111.111', port='1111'))

    # 4. find_all() 测试  --> ok
    # for ip in mongo.find_all():
    #     print(ip)

    # 5. find() 测试
    # 返回 0<speed<1的代理
    print("---find() 测试---")
    li_list = mongo.find(condition={'speed': {"$gte": 0, "$lte": 1}})
    for i in li_list:
        print(i)

    # 6.random_one() 测试
    print("---random_one() 测试---")
    ip = mongo.random_one()
    print(ip)

